%% PDE SOLVER FOR THE OPTIMAL CONTRACT
function [V, C, Beta, K, Eta, f, ExRet, mu_v, mu_y, vol_y, vol_phi, mu_k, mu_c, vol_k, vol_c, Suff] = optimal_system_solver(params, grid_params, num_params, timestamp)
%% PARAMTERS
temp = num2cell(params);
[rho, r, ra, alpha, sigma, lambda, M, eta_min, eta_max] = temp{:};

[n_grid, Y_mat, PHI_mat, dy, dphi] = grid_params{:};

[max_iter, D, thr] = num_params{:};

%% Initialize matrices
V_out = (((ra-r)*(1-rho)/rho)^(rho/(1-rho))); %Value for shutting down the firm
V00 = V_out*ones(n_grid, n_grid);

C0 = 0.01*ones(n_grid,n_grid);
Beta0 = 0.01*ones(n_grid, n_grid);
Eta0 = eta_max*ones(n_grid, n_grid);

%% ITERATIONS
min_iter_lcp = 1;

% surf(Y_mat,PHI_mat,K)

iter = 1;
reiterate = 1;
while reiterate && (iter <= max_iter)
            
    tic        

    %% Find optimal policy.
    [C, Beta, Eta, K, f, ExRet, mu_v, mu_y, vol_y, vol_phi] = optimal_optimal_policy(V00, C0, Beta0, Eta0, params, grid_params);
        
    %% SOLVE VALUE FUNCTION
    V = optimal_hjb_solver(V00,C, Beta, Eta, K, f, ExRet, mu_v, mu_y, vol_y, vol_phi, iter, min_iter_lcp, params, grid_params, num_params);
    
    %% CHECK CHANGE
        
    vmin = min(min(V));
    vmax = max(max(V));
    
    [yposmin, phiposmin] = find(V == vmin,1); 
    [yposmax, phiposmax] = find(V == vmax,1); 
    
    disp(['Min of ', num2str(vmin), ' at ', num2str(yposmin),',',num2str(phiposmin)])
    disp(['Max of ', num2str(vmax,'%.4e'), ' at ', num2str(yposmax),',',num2str(phiposmax)])
    
    disp(['Step size of ', num2str(D,'%.4e')])
        
    changemat = abs(V(1:end,1:end)./V00(1:end,1:end)-1);
    test = max(max(changemat));
    [yposchange, phiposchange] = find(changemat == test,1);
    disp('Optimal Contract')
    disp(['Change of ', num2str(test,'%.4e'), ' at ', num2str(yposchange),',',num2str(phiposchange)])
    
        
    if  test < thr
        reiterate = 0;
        disp(['Converged at iteration ',num2str(iter)])
    elseif sum(sum(isnan(V)))>0
        disp('NaN appeared');
        break
    else
        V00 = V;
        disp(['Iteration ',num2str(iter)])
        C0 = C;
        Eta0 = Eta;
        Beta0 = Beta;
    end
    
    iter = iter+1;
    
    toc
    
    disp(' ')
    disp(' ')
    
end

%% Dynamics of Incentives
Betap = Beta;

%Initialize matrices
Betay_f = zeros(n_grid);
Betay_b = zeros(n_grid);
Betay_c = zeros(n_grid);

Betaphi_f = zeros(n_grid);
Betaphi_b = zeros(n_grid);
Betaphi_c = zeros(n_grid);

Betayphi = zeros(n_grid);

% First Differences w.r.t. y
Betay_f(1:end-1,:) = (Betap(2:end,:) - Betap(1:end-1,:))/dy;
Betay_b(2:end,:) = (Betap(2:end,:) - Betap(1:end-1,:))/dy;
Betay_c(2:end-1,:) = 0.5*(Betay_f(2:end-1,:) + Betay_b(2:end-1,:));

% First Differences w.r.t. phi
Betaphi_f(:,1:end-1) = (Betap(:,2:end) - Betap(:,1:end-1))/dphi;
Betaphi_b(:,2:end) = (Betap(:,2:end) - Betap(:,1:end-1))/dphi;
Betaphi_c(:,2:end-1) = 0.5*(Betaphi_f(:,2:end-1) + Betaphi_b(:,2:end-1));

% VOLATILITY
vol_beta = Betay_c.*vol_y + Betaphi_c.*vol_phi;

%% Dynamics of Capital
Kp = K;

%Initialize matrices
Ky_f = zeros(n_grid);
Ky_b = zeros(n_grid);
Ky_c = zeros(n_grid);

Kphi_f = zeros(n_grid);
Kphi_b = zeros(n_grid);
Kphi_c = zeros(n_grid);

Kyphi = zeros(n_grid);

% First Differences w.r.t. y
Ky_f(1:end-1,:) = (Kp(2:end,:) - Kp(1:end-1,:))/dy;
Ky_b(2:end,:) = (Kp(2:end,:) - Kp(1:end-1,:))/dy;
Ky_c(2:end-1,:) = 0.5*(Ky_f(2:end-1,:) + Ky_b(2:end-1,:));

% First Differences w.r.t. phi
Kphi_f(:,1:end-1) = (Kp(:,2:end) - Kp(:,1:end-1))/dphi;
Kphi_b(:,2:end) = (Kp(:,2:end) - Kp(:,1:end-1))/dphi;
Kphi_c(:,2:end-1) = 0.5*(Kphi_f(:,2:end-1) + Kphi_b(:,2:end-1));

%Second Differences w.r.t. y
Kyy = (Ky_f - Ky_b)/dy;

%Second Differences w.r.t. phi
Kphiphi = (Kphi_f - Kphi_b)/dphi;

%Cross Derivative
Kyphi(2:end-1, 2:end-1) = (2*Kp(2:end-1, 2:end-1) + Kp(3:end, 3:end) + Kp(1:end-2, 1:end-2) +...
         -Kp(3:end, 2:end-1) - Kp(1:end-2, 2:end-1) - Kp(2:end-1, 3:end) - Kp(2:end-1, 1:end-2))/(2*dphi*dy);

% DRIFT
mu_k = Ky_c.*mu_y + 0.5*Kyy.*vol_y.^2 + 0.5*Kphiphi.*vol_phi.^2 + Kyphi.*vol_y.*vol_phi;

% VOLATILITY
vol_k = Ky_c.*vol_y + Kphi_c.*vol_phi;

%% Dynamics of COMPENSATION
Cp = C;

%Initialize matrices
Cy_f = zeros(n_grid);
Cy_b = zeros(n_grid);
Cy_c = zeros(n_grid);

Cphi_f = zeros(n_grid);
Cphi_b = zeros(n_grid);
Cphi_c = zeros(n_grid);

Cyphi = zeros(n_grid);

% First Differences w.r.t. y
Cy_f(1:end-1,:) = (Cp(2:end,:) - Cp(1:end-1,:))/dy;
Cy_b(2:end,:) = (Cp(2:end,:) - Cp(1:end-1,:))/dy;
Cy_c(2:end-1,:) = 0.5*(Cy_f(2:end-1,:) + Cy_b(2:end-1,:));

% First Differences w.r.t. phi
Cphi_f(:,1:end-1) = (Cp(:,2:end) - Cp(:,1:end-1))/dphi;
Cphi_b(:,2:end) = (Cp(:,2:end) - Cp(:,1:end-1))/dphi;
Cphi_c(:,2:end-1) = 0.5*(Cphi_f(:,2:end-1) + Cphi_b(:,2:end-1));

%Second Differences w.r.t. y
Cyy = (Cy_f - Cy_b)/dy;

%Second Differences w.r.t. phi
Cphiphi = (Cphi_f - Cphi_b)/dphi;

%Cross Derivative
Cyphi(2:end-1, 2:end-1) = (2*Cp(2:end-1, 2:end-1) + Cp(3:end, 3:end) + Cp(1:end-2, 1:end-2) +...
         -Cp(3:end, 2:end-1) - Cp(1:end-2, 2:end-1) - Cp(2:end-1, 3:end) - Cp(2:end-1, 1:end-2))/(2*dphi*dy);

% DRIFT
mu_c = Cy_c.*mu_y + 0.5*Cyy.*vol_y.^2 + 0.5*Cphiphi.*vol_phi.^2 + Cyphi.*vol_y.*vol_phi;

% VOLATILITY
vol_c = Cy_c.*vol_y + Cphi_c.*vol_phi;

%% SUFFICIENT CONDITION
% Holds if Suff matrix is positive

V1 = [V(1,:);V;V(end,:)];
V2 = [V1(:,1),V1, V1(:,end)];

%Initialize matrices
Vy_f = zeros(n_grid+2,n_grid+2);
Vy_b = zeros(n_grid+2,n_grid+2);
Vy_c = zeros(n_grid+2, n_grid+2);

Vphi_f = zeros(n_grid+2,n_grid+2);
Vphi_b = zeros(n_grid+2,n_grid+2);
Vphi_c = zeros(n_grid+2, n_grid+2);

Vyphi = zeros(n_grid+2,n_grid+2);

% First Differences w.r.t. y
Vy_f(1:end-1,:) = (V2(2:end,:) - V2(1:end-1,:))/dy;
Vy_b(2:end,:) = (V2(2:end,:) - V2(1:end-1,:))/dy;
Vy_c(2:end-1,:) = 0.5*(Vy_f(2:end-1,:) + Vy_b(2:end-1,:));

%Vy_f(end,:) = 10^(60);


% First Differences w.r.t. phi
Vphi_f(:,1:end-1) = (V2(:,2:end) - V2(:,1:end-1))/dphi;
Vphi_b(:,2:end) = (V2(:,2:end) - V2(:,1:end-1))/dphi;
Vphi_c(:,2:end-1) = 0.5*(Vphi_f(:,2:end-1) + Vphi_b(:,2:end-1));

%Second Differences w.r.t. y
Vyy = (Vy_f - Vy_b)/dy;

%Second Differences w.r.t. phi
Vphiphi = (Vphi_f - Vphi_b)/dphi;

%Cross Derivative           
Vyphi(:, 2:end-1) = 0.5*((Vy_f(:,3:end) - Vy_f(:,2:end-1))/dphi + (Vy_f(:,2:end-1) - Vy_f(:,1:end-2))/dphi);

% Reduce to original grid
Vy_b = Vy_b(2:end-1,2:end-1);
Vy_b(end,:) = 0;
Vy_f = Vy_f(2:end-1,2:end-1);
Vy_f(end,:) = 0;
Vy_c = Vy_c(2:end-1,2:end-1);
Vy_c(end,:) = 0;
Vphi_c = Vphi_c(2:end-1,2:end-1);
Vyy = Vyy(2:end-1,2:end-1); 
Vyy(end,:) = 0;
Vphiphi = Vphiphi(2:end-1,2:end-1); 
Vyphi = Vyphi(2:end-1,2:end-1);
Vyphi(end,:) = 0;

L = ( ((1-rho)*Vy_f - rho*Y_mat.*Vyy).*Beta.*(1-PHI_mat) + ((1-PHI_mat).*Vyphi - Y_mat.*Vyy - Vy_f).*Eta.*PHI_mat.*(1-PHI_mat) );
R = ( Eta.*(1-2*PHI_mat).*(1-PHI_mat).*Vy_f );

% Difference between vol of info rent and (1-2phi)*eta*InfoRent (all scaled by v^(1-rho))
Suff = L - R;

end
